#!/usr/bin/env python3
"""
将不同程序（OpenSSL、WolfSSL等）生成的JSON文件转换为统一的TLS输入输出状态机
使用大模型进行语义对齐，支持跨平台状态机转换
"""

import re
import json
import argparse
import os
import asyncio
import subprocess
from typing import Dict, List, Tuple, Set, Optional
from collections import defaultdict
import time
import utils


class LLMSemanticMapper:
    """基于大模型的语义对齐映射器"""
    
    def __init__(self, use_async: bool = False):
        self.use_async = use_async
        self.cache = {}  # 缓存映射结果
        
    def call_llm(self, prompt: str, max_retries: int = 3) -> str:
        """调用大模型API"""
        for attempt in range(max_retries):
            try:
                if self.use_async:
                    result = self._run_async_llm_call(prompt)
                    if result and isinstance(result, dict):
                        return result.get("result", "ERROR: 无效响应")
                    else:
                        return "ERROR: 无效响应"
                else:
                    result = utils.get_ans_json(prompt)
                    return result.get("result", "ERROR: 无效响应")
            except Exception as e:
                print(f"LLM API调用失败 (尝试 {attempt + 1}/{max_retries}): {e}")
                if attempt < max_retries - 1:
                    # 增加延迟时间，避免过于频繁的重试
                    delay = (2 ** attempt) + 1  # 1, 3, 7秒
                    time.sleep(delay)
                else:
                    return "ERROR: LLM调用失败"
    
    def map_state_to_tls_phase(self, state: str, context: str = "") -> str:
        """将状态映射到TLS握手阶段"""
        cache_key = f"state_{state}_{context}"
        if cache_key in self.cache:
            return self.cache[cache_key]
        
        prompt = f"""
你是一个TLS协议专家。请分析以下状态名称，将其映射到标准的TLS 1.3握手阶段。

状态名称: {state}
上下文: {context}

请从以下TLS握手阶段中选择最匹配的一个：
1. CLIENT_HELLO - 客户端发送Hello消息
2. SERVER_HELLO - 服务器响应Hello消息  
3. ENCRYPTED_EXTENSIONS - 加密扩展交换
4. CERTIFICATE - 证书交换
5. CERTIFICATE_VERIFY - 证书验证
6. FINISHED - 握手完成
7. APPLICATION_DATA - 应用数据传输
8. KEY_UPDATE - 密钥更新
9. NEW_SESSION_TICKET - 新会话票据
10. HELLO_RETRY_REQUEST - Hello重试请求
11. END_OF_EARLY_DATA - 早期数据结束
12. UNKNOWN - 无法确定

请返回JSON格式：{{"result": "阶段名称"}}
"""
        
        result = self.call_llm(prompt)
        self.cache[cache_key] = result
        return result
    
    def map_condition_to_tls_message(self, condition: str, from_state: str, to_state: str) -> Tuple[str, str]:
        """将条件映射到TLS消息类型"""
        cache_key = f"condition_{condition}_{from_state}_{to_state}"
        if cache_key in self.cache:
            return self.cache[cache_key]
        
        # 首先尝试从条件中直接提取消息类型
        message_type = self._extract_message_type_from_condition(condition)
        
        if message_type != "Unknown":
            # 如果能从条件中提取到消息类型，需要根据方向确定输入输出
            if "Dir: SEND" in condition:
                # 发送消息，输出是消息类型，输入是前一个状态的消息
                output_msg = message_type
                input_msg = "Unknown"  # 需要进一步分析
            elif "Dir: RECEIVE" in condition:
                # 接收消息，输入是消息类型，输出是响应消息
                input_msg = message_type
                output_msg = "Unknown"  # 需要进一步分析
            else:
                # 无法确定方向，使用LLM进一步分析
                input_msg = "Unknown"
                output_msg = "Unknown"
            
            # 如果还有未知的消息，使用LLM进一步分析
            if input_msg == "Unknown" or output_msg == "Unknown":
                # 使用LLM分析完整的输入输出
                result = self._analyze_input_output_with_llm(condition, from_state, to_state, message_type)
                if result:
                    if input_msg == "Unknown":
                        input_msg = result.get("input_message", "Unknown")
                    if output_msg == "Unknown":
                        output_msg = result.get("output_message", "Unknown")
        else:
            # 否则使用LLM进行更详细的分析
            prompt = f"""
你是一个TLS协议专家。请分析以下状态转换条件，确定涉及的TLS消息类型。

转换条件: {condition}
源状态: {from_state}
目标状态: {to_state}

请仔细分析条件中的关键词：
- 如果条件包含 "Msg: XXX"，则消息类型是 XXX
- 如果条件包含 "Dir: SEND"，则这是发送消息
- 如果条件包含 "Dir: RECEIVE"，则这是接收消息
- 如果条件包含 "hello_retry_request"，则消息类型是 HelloRetryRequest
- 如果条件包含 "client_hello"，则消息类型是 ClientHello
- 如果条件包含 "server_hello"，则消息类型是 ServerHello
- 如果条件包含 "certificate"，则消息类型是 Certificate
- 如果条件包含 "finished"，则消息类型是 Finished
- 如果条件包含 "encrypted_extensions"，则消息类型是 EncryptedExtensions

标准TLS 1.3消息类型：
- ClientHello: 客户端Hello消息
- ServerHello: 服务器Hello消息
- EncryptedExtensions: 加密扩展
- Certificate: 证书消息
- CertificateVerify: 证书验证
- Finished: 完成消息
- EndOfEarlyData: 早期数据结束
- KeyUpdate: 密钥更新
- NewSessionTicket: 新会话票据
- HelloRetryRequest: Hello重试请求
- ChangeCipherSpec: 密码套件变更
- ApplicationData: 应用数据
- Alert: 告警消息

请返回JSON格式：{{"input_message": "输入消息类型", "output_message": "输出消息类型"}}

如果无法确定，返回: {{"input_message": "Unknown", "output_message": "Unknown"}}
"""
            
            result = self.call_llm(prompt)
            try:
                if isinstance(result, dict):
                    input_msg = result.get("input_message", "Unknown")
                    output_msg = result.get("output_message", "Unknown")
                else:
                    input_msg, output_msg = "Unknown", "Unknown"
            except:
                input_msg, output_msg = "Unknown", "Unknown"
        
        self.cache[cache_key] = (input_msg, output_msg)
        return (input_msg, output_msg)
    
    def _extract_message_type_from_condition(self, condition: str) -> str:
        """从条件中直接提取消息类型"""
        condition_lower = condition.lower()
        
        # 检查是否包含 "Msg: XXX" 格式
        msg_match = re.search(r'msg:\s*([a-z_]+)', condition_lower)
        if msg_match:
            msg_type = msg_match.group(1).upper()
            # 标准化消息类型名称
            if msg_type == "CLIENT_HELLO":
                return "ClientHello"
            elif msg_type == "SERVER_HELLO":
                return "ServerHello"
            elif msg_type == "ENCRYPTED_EXTENSIONS":
                return "EncryptedExtensions"
            elif msg_type == "CERTIFICATE":
                return "Certificate"
            elif msg_type == "CERTIFICATE_VERIFY":
                return "CertificateVerify"
            elif msg_type == "FINISHED":
                return "Finished"
            elif msg_type == "HELLO_RETRY_REQUEST":
                return "HelloRetryRequest"
            elif msg_type == "CHANGE_CIPHER_SPEC":
                return "ChangeCipherSpec"
            else:
                return msg_type
        
        # 检查其他关键词
        if "hello_retry_request" in condition_lower:
            return "HelloRetryRequest"
        elif "client_hello" in condition_lower:
            return "ClientHello"
        elif "server_hello" in condition_lower:
            return "ServerHello"
        elif "certificate" in condition_lower:
            return "Certificate"
        elif "finished" in condition_lower:
            return "Finished"
        elif "encrypted_extensions" in condition_lower:
            return "EncryptedExtensions"
        elif "change_cipher_spec" in condition_lower:
            return "ChangeCipherSpec"
        
        return "Unknown"
    
    def _analyze_input_output_with_llm(self, condition: str, from_state: str, to_state: str, known_message: str) -> Optional[Dict]:
        """使用LLM分析输入输出消息"""
        prompt = f"""
你是一个TLS协议专家。请分析以下状态转换，确定输入和输出消息类型。

转换条件: {condition}
源状态: {from_state}
目标状态: {to_state}
已知消息类型: {known_message}

在TLS握手过程中，输入和输出通常是不同的消息。请根据TLS协议规范分析：

1. 如果这是发送消息（Dir: SEND），那么：
   - 输出消息是已知的消息类型
   - 输入消息通常是前一个状态接收到的消息

2. 如果这是接收消息（Dir: RECEIVE），那么：
   - 输入消息是已知的消息类型
   - 输出消息通常是响应消息

3. 考虑TLS握手的典型流程：
   - ClientHello -> ServerHello
   - ServerHello -> EncryptedExtensions
   - EncryptedExtensions -> Certificate
   - Certificate -> CertificateVerify
   - CertificateVerify -> Finished
   - Finished -> ApplicationData

请返回JSON格式：{{"input_message": "输入消息类型", "output_message": "输出消息类型"}}
"""
        
        try:
            if self.use_async:
                # 异步调用 - 使用新的方法
                result = self._run_async_llm_call(prompt)
            else:
                # 同步调用
                result = utils.get_ans_json(prompt)
            
            if result:
                return result
        except Exception as e:
            print(f"LLM分析输入输出失败: {e}")
        
        return None
    
    def _run_async_llm_call(self, prompt: str) -> str:
        """运行异步LLM调用的辅助方法"""
        import asyncio
        import threading
        
        def run_in_thread():
            """在新线程中运行异步代码"""
            try:
                return asyncio.run(utils.get_ans_json_async(prompt))
            except Exception as e:
                print(f"异步LLM调用失败: {e}")
                return None
        
        # 在新线程中运行异步代码，避免事件循环冲突
        result = None
        thread = threading.Thread(target=lambda: setattr(self, '_async_result', run_in_thread()))
        thread.start()
        thread.join()
        
        return getattr(self, '_async_result', None)
    
    def determine_direction(self, state: str) -> str:
        """确定状态的输入输出方向"""
        cache_key = f"direction_{state}"
        if cache_key in self.cache:
            return self.cache[cache_key]
        
        prompt = f"""
分析以下TLS状态名称，确定它是输入状态还是输出状态：

状态名称: {state}

规则：
- 如果状态名称包含 "CW_" (Client Write) 或 "SW_" (Server Write)，则为输出状态
- 如果状态名称包含 "CR_" (Client Read) 或 "SR_" (Server Read)，则为输入状态
- 如果状态名称包含 "SEND"、"WRITE"、"OUTPUT" 等，则为输出状态
- 如果状态名称包含 "RECEIVE"、"READ"、"INPUT" 等，则为输入状态
- 如果状态名称包含 "COMPLETE"、"DONE"、"FINISHED" 等，需要根据上下文判断

请返回JSON格式：{{"result": "IN"}} 或 {{"result": "OUT"}}
"""
        
        result = self.call_llm(prompt)
        self.cache[cache_key] = result
        return result


def generate_pdf_from_dot(dot_file_path: str, output_dir: str = None) -> str:
    """
    从DOT文件生成PDF文件
    
    Args:
        dot_file_path: DOT文件路径
        output_dir: 输出目录，如果为None则使用DOT文件所在目录
    
    Returns:
        str: 生成的PDF文件路径
    """
    try:
        # 检查graphviz是否可用
        result = subprocess.run(['dot', '-V'], capture_output=True, text=True)
        if result.returncode != 0:
            print("警告: graphviz未安装或不在PATH中，无法生成PDF文件")
            return None
    except FileNotFoundError:
        print("警告: graphviz未安装或不在PATH中，无法生成PDF文件")
        return None
    
    # 确定输出目录和文件名
    if output_dir is None:
        output_dir = os.path.dirname(dot_file_path)
    
    base_name = os.path.splitext(os.path.basename(dot_file_path))[0]
    pdf_file_path = os.path.join(output_dir, f"{base_name}.pdf")
    
    try:
        # 使用dot命令生成PDF
        command = ['dot', '-Tpdf', dot_file_path, '-o', pdf_file_path]
        result = subprocess.run(command, capture_output=True, text=True)
        
        if result.returncode == 0:
            print(f"成功生成PDF文件: {pdf_file_path}")
            return pdf_file_path
        else:
            print(f"生成PDF失败: {result.stderr}")
            return None
            
    except Exception as e:
        print(f"生成PDF时出错: {e}")
        return None


class UniversalJSONParser:
    """通用JSON文件解析器"""
    
    def __init__(self):
        self.states: Set[str] = set()
        self.transitions: List[Tuple[str, str, str]] = []
        self.state_descriptions: Dict[str, str] = {}
    
    def parse_json_file(self, json_file: str):
        """解析JSON文件"""
        with open(json_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        self._extract_states_and_transitions(data)
    
    def _extract_states_and_transitions(self, data: Dict):
        """从JSON数据中提取状态和转换"""
        # 遍历JSON中的所有顶级键（如"server", "client", "server1.3", "client1.3"等）
        for protocol_key, protocol_data in data.items():
            if isinstance(protocol_data, dict):
                # 检查是否是新的格式（直接包含states和transitions）
                if 'states' in protocol_data and 'transitions' in protocol_data:
                    # 标准格式：直接包含states和transitions
                    for state in protocol_data['states']:
                        if isinstance(state, str) and state.strip():
                            self.states.add(state.strip())
                    
                    for transition in protocol_data['transitions']:
                        self._process_transition(transition)
                
                else:
                    # 新格式：包含函数，每个函数有自己的transitions
                    for func_name, func_data in protocol_data.items():
                        if isinstance(func_data, dict) and 'transitions' in func_data:
                            for transition in func_data['transitions']:
                                self._process_transition(transition)
    
    def _process_transition(self, transition: Dict):
        """处理单个转换"""
        if isinstance(transition, dict):
            start_state = transition.get('start_state', '').strip()
            end_state = transition.get('end_state', '').strip()
            condition = transition.get('condition', '').strip()
            
            if start_state and end_state:
                self.states.add(start_state)
                self.states.add(end_state)
                
                # 构建条件字符串，包含消息类型和方向信息
                condition_parts = []
                if condition:
                    condition_parts.append(f"Cond: {condition}")
                
                message_type = transition.get('message_type', '').strip()
                if message_type:
                    condition_parts.append(f"Msg: {message_type}")
                
                message_direction = transition.get('message_direction', '').strip()
                if message_direction:
                    condition_parts.append(f"Dir: {message_direction}")
                
                condition_explanation = transition.get('condition explanation', '').strip()
                if condition_explanation:
                    condition_parts.append(f"Explanation: {condition_explanation}")
                
                full_condition = '\n'.join(condition_parts)
                self.transitions.append((start_state, end_state, full_condition))
    
    def get_states(self) -> Set[str]:
        return self.states
    
    def get_transitions(self) -> List[Tuple[str, str, str]]:
        return self.transitions
    
    def get_state_description(self, state: str) -> str:
        return self.state_descriptions.get(state, "")


class UniversalIOStateMachine:
    """通用输入输出状态机"""
    
    def __init__(self, llm_mapper: LLMSemanticMapper):
        self.states: Set[str] = set()
        self.io_transitions: Dict[str, List[Tuple[str, str, str]]] = defaultdict(list)
        self.llm_mapper = llm_mapper
        self.state_phases: Dict[str, str] = {}
    
    def add_state(self, state: str):
        """添加状态"""
        self.states.add(state)
        if state not in self.io_transitions:
            self.io_transitions[state] = []
    
    def add_io_transition(self, from_state: str, to_state: str, input_msg: str, output_msg: str):
        """添加输入输出转换"""
        self.add_state(from_state)
        self.add_state(to_state)
        self.io_transitions[from_state].append((input_msg, output_msg, to_state))
    
    def convert_from_parser(self, parser_obj):
        """从解析器转换状态机"""
        print("开始语义对齐转换...")
        
        if self.llm_mapper.use_async:
            # 异步并发处理
            self._convert_async(parser_obj)
        else:
            # 同步顺序处理
            self._convert_sync(parser_obj)
    
    def convert_single_protocol(self, protocol_name, protocol_data):
        """转换单个协议的状态机"""
        print(f"处理协议: {protocol_name}")
        
        # 创建新的状态机实例
        protocol_fsm = UniversalIOStateMachine(self.llm_mapper)
        
        # 收集所有状态和转换
        all_states = set()
        all_transitions = []
        
        for func_name, func_data in protocol_data.items():
            if isinstance(func_data, dict) and 'transitions' in func_data:
                for transition in func_data['transitions']:
                    start_state = transition.get('start_state')
                    end_state = transition.get('end_state')
                    condition = transition.get('condition', '')
                    
                    if start_state and end_state:
                        all_states.add(start_state)
                        all_states.add(end_state)
                        all_transitions.append((start_state, end_state, condition))
        
        # 处理状态映射
        for state in all_states:
            phase = self.llm_mapper.map_state_to_tls_phase(state, "")
            protocol_fsm.state_phases[state] = phase
            protocol_fsm.states.add(state)
            print(f"  状态 {state} -> TLS阶段: {phase}")
        
        # 处理转换
        for i, (start_state, end_state, condition) in enumerate(all_transitions):
            print(f"  处理转换 {i+1}/{len(all_transitions)}: {start_state} -> {end_state}")
            protocol_fsm._process_transition(start_state, end_state, condition)
        
        return protocol_fsm
    
    def _convert_sync(self, parser_obj):
        """同步转换处理"""
        # 首先确定所有状态的TLS阶段
        for state in parser_obj.get_states():
            description = parser_obj.get_state_description(state)
            phase = self.llm_mapper.map_state_to_tls_phase(state, description)
            self.state_phases[state] = phase
            print(f"状态 {state} -> TLS阶段: {phase}")
        
        # 然后处理转换
        for i, (start_state, end_state, condition) in enumerate(parser_obj.get_transitions()):
            print(f"处理转换 {i+1}/{len(parser_obj.get_transitions())}: {start_state} -> {end_state}")
            self._process_transition(start_state, end_state, condition)
    
    def _convert_async(self, parser_obj):
        """异步并发转换处理"""
        import asyncio
        import concurrent.futures
        import time
        
        # 首先确定所有状态的TLS阶段
        print("处理状态映射...")
        with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:  # 减少并发数
            state_futures = {}
            for state in parser_obj.get_states():
                description = parser_obj.get_state_description(state)
                future = executor.submit(self._process_state_with_delay, state, description)
                state_futures[state] = future
            
            # 收集结果
            for state, future in state_futures.items():
                try:
                    phase = future.result()
                    self.state_phases[state] = phase
                    print(f"状态 {state} -> TLS阶段: {phase}")
                except Exception as e:
                    print(f"状态 {state} 处理失败: {e}")
                    self.state_phases[state] = "UNKNOWN"
        
        # 然后并发处理转换
        print("处理转换映射...")
        with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:  # 减少并发数
            transition_futures = {}
            for i, (start_state, end_state, condition) in enumerate(parser_obj.get_transitions()):
                future = executor.submit(self._process_transition_with_delay, start_state, end_state, condition, i)
                transition_futures[i] = (future, start_state, end_state)
            
            # 收集结果
            for i, (future, start_state, end_state) in transition_futures.items():
                try:
                    future.result()
                    print(f"处理转换 {i+1}/{len(parser_obj.get_transitions())}: {start_state} -> {end_state}")
                except Exception as e:
                    print(f"处理转换 {i+1} 失败: {e}")
    
    def _process_state_with_delay(self, state, description):
        """处理状态映射，添加延迟"""
        import time
        time.sleep(0.1)  # 添加100ms延迟
        return self.llm_mapper.map_state_to_tls_phase(state, description)
    
    def _process_transition_with_delay(self, start_state, end_state, condition, index):
        """处理转换映射，添加延迟"""
        import time
        time.sleep(0.2)  # 添加200ms延迟
        return self._process_transition(start_state, end_state, condition)
    
    def _process_transition(self, start_state, end_state, condition):
        """处理单个转换"""
        # 使用LLM映射条件到TLS消息
        input_msg, output_msg = self.llm_mapper.map_condition_to_tls_message(
            condition, start_state, end_state
        )
        
        # 确定方向
        start_direction = self.llm_mapper.determine_direction(start_state)
        end_direction = self.llm_mapper.determine_direction(end_state)
        
        # 构建输入输出标签
        if start_direction == "OUT" and end_direction == "IN":
            # 发送方到接收方
            actual_input = f"IN:{input_msg}" if input_msg != "Unknown" else "IN:Message"
            actual_output = f"OUT:{output_msg}" if output_msg != "Unknown" else "OUT:Message"
        elif start_direction == "IN" and end_direction == "OUT":
            # 接收方到发送方
            actual_input = f"IN:{input_msg}" if input_msg != "Unknown" else "IN:Message"
            actual_output = f"OUT:{output_msg}" if output_msg != "Unknown" else "OUT:Message"
        else:
            # 其他情况
            actual_input = f"IN:{input_msg}" if input_msg != "Unknown" else "IN:Message"
            actual_output = f"OUT:{output_msg}" if output_msg != "Unknown" else "OUT:Message"
        
        self.add_io_transition(start_state, end_state, actual_input, actual_output)
    
    def _clean_dot_label(self, label: str) -> str:
        """清理DOT标签中的特殊字符"""
        if not label:
            return "Unknown"
        
        # 替换特殊字符
        cleaned = label.replace('"', '\\"')  # 转义双引号
        cleaned = cleaned.replace('\n', '\\n')  # 转义换行符
        cleaned = cleaned.replace('\r', '\\r')  # 转义回车符
        cleaned = cleaned.replace('\t', '\\t')  # 转义制表符
        cleaned = cleaned.replace('\\', '\\\\')  # 转义反斜杠
        
        # 限制长度，避免过长的标签
        if len(cleaned) > 50:
            cleaned = cleaned[:47] + "..."
        
        return cleaned
    
    def to_dot(self) -> str:
        """转换为DOT格式"""
        dot_lines = ["digraph Universal_IO_FSM {"]
        dot_lines.append("    rankdir=LR;")
        dot_lines.append("    node [shape=box, style=rounded, fontsize=10];")
        dot_lines.append("    edge [fontsize=8];")
        dot_lines.append("")
        
        # 简化：不按阶段分组，直接列出所有状态
        for state in sorted(self.states):
            # 简化状态标签
            short_label = state.replace("TLS_ST_", "").replace("DTLS_ST_", "")
            dot_lines.append(f'    {state} [label="{short_label}"];')
        
        dot_lines.append("")
        
        # 添加转换，简化标签
        for from_state, transitions in self.io_transitions.items():
            for input_msg, output_msg, to_state in transitions:
                # 简化标签，只保留关键信息
                clean_input = self._clean_dot_label(input_msg)
                clean_output = self._clean_dot_label(output_msg)
                
                # 进一步简化标签
                if "IN:" in clean_input and "OUT:" in clean_output:
                    input_part = clean_input.replace("IN:", "").split("/")[0]
                    output_part = clean_output.replace("OUT:", "").split("/")[0]
                    label = f"{input_part}/{output_part}"
                else:
                    label = f"{clean_input}/{clean_output}"
                
                # 限制标签长度
                if len(label) > 30:
                    label = label[:27] + "..."
                
                dot_lines.append(f'    {from_state} -> {to_state} [label="{label}"];')
        
        dot_lines.append("}")
        return "\n".join(dot_lines)
    
    def to_json(self) -> Dict:
        """转换为JSON格式"""
        result = {
            "metadata": {
                "type": "Universal_TLS_IO_StateMachine",
                "total_states": len(self.states),
                "total_transitions": sum(len(transitions) for transitions in self.io_transitions.values())
            },
            "states": {},
            "transitions": []
        }
        
        # 添加状态信息
        for state in self.states:
            phase = self.state_phases.get(state, "UNKNOWN")
            result["states"][state] = {
                "tls_phase": phase,
                "outgoing_transitions": len(self.io_transitions.get(state, []))
            }
        
        # 添加转换信息
        for from_state, transitions in self.io_transitions.items():
            for input_msg, output_msg, to_state in transitions:
                result["transitions"].append({
                    "from_state": from_state,
                    "to_state": to_state,
                    "input": input_msg,
                    "output": output_msg,
                    "from_phase": self.state_phases.get(from_state, "UNKNOWN"),
                    "to_phase": self.state_phases.get(to_state, "UNKNOWN")
                })
        
        return result


def main():
    parser = argparse.ArgumentParser(description='通用JSON到TLS输入输出状态机转换器')
    parser.add_argument('-i', '--input', required=True, help='输入JSON文件路径')
    parser.add_argument('-o', '--output', required=True, help='输出目录')
    parser.add_argument('--async', action='store_true', help='使用异步LLM调用')
    parser.add_argument('--cache-file', help='缓存文件路径（可选）')
    
    args = parser.parse_args()
    
    # 创建输出目录
    os.makedirs(args.output, exist_ok=True)
    
    # 初始化LLM映射器
    llm_mapper = LLMSemanticMapper(use_async=getattr(args, 'async', False))
    
    # 加载缓存（如果存在）
    if args.cache_file and os.path.exists(args.cache_file):
        try:
            with open(args.cache_file, 'r', encoding='utf-8') as f:
                llm_mapper.cache = json.load(f)
            print(f"加载缓存: {len(llm_mapper.cache)} 个条目")
        except Exception as e:
            print(f"加载缓存失败: {e}")
    
    # 解析JSON文件
    print(f"解析JSON文件: {args.input}")
    with open(args.input, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    # 为每个协议分别生成状态机
    base_name = os.path.splitext(os.path.basename(args.input))[0]
    
    for protocol_name, protocol_data in data.items():
        if not isinstance(protocol_data, dict):
            continue
            
        print(f"\n=== 处理协议: {protocol_name} ===")
        
        # 创建状态机实例
        io_fsm = UniversalIOStateMachine(llm_mapper)
        
        # 转换单个协议
        io_fsm = io_fsm.convert_single_protocol(protocol_name, protocol_data)
        
        # 生成输出文件
        protocol_base = f"{base_name}_{protocol_name}_io"
        
        # 生成DOT文件
        dot_output = os.path.join(args.output, f"{protocol_base}.dot")
        with open(dot_output, 'w', encoding='utf-8') as f:
            f.write(io_fsm.to_dot())
        print(f"生成DOT文件: {dot_output}")
        
        # 生成JSON文件
        json_output = os.path.join(args.output, f"{protocol_base}.json")
        with open(json_output, 'w', encoding='utf-8') as f:
            json.dump(io_fsm.to_json(), f, indent=2, ensure_ascii=False)
        print(f"生成JSON文件: {json_output}")
        
        # 生成PDF文件
        print(f"正在生成 {protocol_name} 的PDF文件...")
        pdf_path = generate_pdf_from_dot(dot_output, args.output)
        if pdf_path:
            print(f"PDF文件已生成: {pdf_path}")
        else:
            print(f"{protocol_name} PDF文件生成失败。")
    
    # 保存缓存
    if args.cache_file:
        try:
            with open(args.cache_file, 'w', encoding='utf-8') as f:
                json.dump(llm_mapper.cache, f, indent=2, ensure_ascii=False)
            print(f"\n保存缓存: {len(llm_mapper.cache)} 个条目")
        except Exception as e:
            print(f"保存缓存失败: {e}")
    
    print("\n所有协议转换完成!")

# python dot_to_io_fsm_llm.py -i save/wolfssl_3.4.1/dots/server1.3.dot -o output --async --cache-file cache.json

if __name__ == '__main__':
    main()
